package datastruct;

import java.util.Objects;

/**
 * 前缀树，也称为字典树或者单词查找树
 * 26个小写字母的前缀树 该段代码示例 参考《剑指offer 突击专项版》第10章节
 *
 * @author 胡宇轩
 * @Email: yuxuan.hu01@bianlifeng.com
 */
public class Trie {
    TrieNode root;

    public Trie(){
        root = new TrieNode();
    }
    /**
     * 前缀树
     * */
    class TrieNode{
        boolean isWord;
        TrieNode[] children;

        public TrieNode(){
            children = new TrieNode[26];
        }
    }

    /**
     * 插入一个单词到前缀树
     * */
    public void insert(String str){
        if(Objects.isNull(str) || str.equals("")){
            return;
        }
        insertStr(str);
    }

    private void insertStr(String str){
        TrieNode rNode = root;
        for(char c : str.toCharArray()){
            if(rNode.children[c - 'a'] == null){
                rNode.children[c - 'a'] = new TrieNode();
            }
            rNode = rNode.children[c - 'a'];
        }
        rNode.isWord = true;
    }

    /**
     * 查找单词是否存在
     * */
    public boolean search(String str){
        TrieNode rNode = root;
        for(char c : str.toCharArray()){
            if(rNode.children[c - 'a'] == null){
                return false;
            }
            rNode = rNode.children[c - 'a'];
        }
        return rNode.isWord;
    }

    /**
     * 查找是否存在以字符串开头的单词
     * */
    public boolean startWith(String prefix){
        TrieNode rNode = root;
        for(char c : prefix.toCharArray()){
            if(rNode.children[c - 'a'] == null){
                return false;
            }
            rNode = rNode.children[c - 'a'];
        }
        return true;
    }
}
